import cv2 
import time
import os, sys
import os.path as osp
from datetime import datetime
import numpy as np


def save_img_by_rtsp_history(save_dir, cap_ip_info, 
    start_time="2021-05-28 15:30:00", 
    end_time="2021-08-26 15:30:00",
    intervel=60 * 60):
    
    cap_ip, user, passwd, channels = cap_ip_info

    #
    for channel in range(channels):
        channel_ip = cap_ip
        print("channel_ip: ", channel_ip)

        start = np.int64(time.mktime(time.strptime(start_time, '%Y-%m-%d %H:%M:%S')))
        end = np.int64(time.mktime(time.strptime(end_time, '%Y-%m-%d %H:%M:%S')))

        continue_cnt = 0
        for ti in range(start, end, intervel):
            # 跳过附近的帧
            if continue_cnt > 0:
                continue_cnt -= 1
                continue
            y_m_d = time.strftime("%Y-%m-%d", time.gmtime(ti + 60*60*8)).split("-")
            sub_path = os.sep.join(y_m_d)
            print("sub_path: ", sub_path)
            cur_save_dir = osp.join(save_dir, sub_path, channel_ip)
            print("cur_save_dir: ", cur_save_dir)
            os.makedirs(cur_save_dir, exist_ok=True)
            img_path = osp.join(cur_save_dir, "{}.jpg".format(ti))
            if osp.exists(img_path): 
                print("exists: ", img_path)
                continue

            cur_time = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime(ti + 60*60*8))
            rtsp = "rtsp://{}:{}@{}:554/Streaming/tracks/{}{}?starttime={}".format(user, passwd, cap_ip, channel+1, "01", cur_time)
            # rtsp = "rtsp://{}:{}@{}:554/Streaming/Channels/{}{}?starttime={}".format(user, passwd, cap_ip, channel+1, "01", cur_time)
            # rtsp = "rtsp://{}:{}@{}:554/Streaming/Channels/{}{}".format(user, passwd, cap_ip, channel+1, "01")
            # rtsp = "rtsp://{}:{}@{}:554/h264/ch{}/sub/av_stream?starttime={}".format(user, passwd, cap_ip, channel+33, cur_time)

            print(rtsp)
            cap = cv2.VideoCapture(rtsp)
            ret,frame = cap.read()
            if not ret: 
                print("ret: ", ret)
                continue_cnt = 20
                continue

            cv2.imwrite(img_path, frame)
            cap.release()  

def save_video_by_rtsp_history(save_dir, cap_ip_info, 
    start_time="2021-05-28 15:30:00", 
    end_time="2021-08-26 15:30:00",
    intervel=60 * 60):
    
    cap_ip, user, passwd, channels = cap_ip_info

    for channel in range(channels):
        channel_ip = ".".join(cap_ip.split(".")[:-1]) +"."+ str(int(cap_ip.split(".")[-1]) - (channels - channel))
        print("channel_ip: ", channel_ip)

        start = np.int64(time.mktime(time.strptime(start_time, '%Y-%m-%d %H:%M:%S')))
        end = np.int64(time.mktime(time.strptime(end_time, '%Y-%m-%d %H:%M:%S')))

        for ti in range(start, end, intervel):

            y_m_d = time.strftime("%Y-%m-%d", time.gmtime(ti + 60*60*8)).split("-")
            sub_path = os.sep.join(y_m_d)
            print("sub_path: ", sub_path)
            cur_save_dir = osp.join(save_dir, sub_path, channel_ip)
            print("cur_save_dir: ", cur_save_dir)
            os.makedirs(cur_save_dir, exist_ok=True)
            fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
            video_path = osp.join(cur_save_dir, "{}.mp4".format(ti))
            if osp.exists(video_path): 
                print("exists: ", video_path)
                continue

            save_strip = time.strftime("%Y-%m-%d_%H-%M-%S", time.gmtime(ti + 60*60*8))
            cur_time = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime(ti + 60*60*8))
            rtsp = "rtsp://{}:{}@{}:554/Streaming/tracks/{}/?starttime={}".format(user, passwd, cap_ip, 101 + channel, cur_time)
            cap = cv2.VideoCapture(rtsp)
            ret,frame = cap.read()
            if not ret: continue

            width = cap.get(3)
            height = cap.get(4)
            fps = 20
            print("width: " , width, " height: ", height, " fps: ", cap.get(5))

            writer = cv2.VideoWriter(video_path, fourcc, fps, (int(width), int(height)))
            frame_cnt = 0
            while (frame_cnt < 10 * fps):
                frame_cnt += 1
                ret,frame = cap.read()
                if not ret:
                    print("frame is empty")
                    break
                writer.write(frame)
            cap.release()  
            writer.release()

def save_img_by_rtsp(save_dir, cap_ip_info):
    cap_ip, user, passwd = cap_ip_info
    cap = cv2.VideoCapture("rtsp://{}:{}@{}:554/ch1/sub/av_stream".format(user, passwd, cap_ip))
    #    const std::string videoStreamAddress = "rtsp://admin:swlskjyxgs0617@192.168.31.111:554/h264/ch34/main/av_stream";
    print("ok")
    ret,frame = cap.read()
    print("ret: ", ret, cap_ip)
    if not ret: 
        return ret

    cur_save_dir = osp.join(save_dir, datetime.now().strftime("%Y-%m-%d"), cap_ip)
    os.makedirs(cur_save_dir, exist_ok=True)
    height, width = frame.shape[:2]
    
    ret,frame = cap.read()
    if not ret:
        print("frame is empty")
        return ret
    cv2.imwrite(osp.join(cur_save_dir, "{}.jpg".format(datetime.now().strftime("%Y-%m-%d %H-%M-%S"))), frame)
    cap.release()  
    return True 

def save_video_by_rtsp(save_dir, cap_ip, save_time):
    cap_ip, user, passwd = cap_ip
    """rtsp格式"""
    cap = cv2.VideoCapture("rtsp://{}:{}@{}:554/ch1/sub/av_stream".format(user, passwd, cap_ip))
    print("ok")
    ret,frame = cap.read()
    print("ret: ", ret)
    if not ret:
        return ret
    # 更新时间
    ti = time.time()
    y_m_d = time.strftime("%Y-%m-%d", time.gmtime(ti)).split("-")
    sub_path = os.sep.join(y_m_d)
    print("sub_path: ", sub_path)
    cur_save_dir = osp.join(save_dir, sub_path, cap_ip)
    print("cur_save_dir: ", cur_save_dir)
    os.makedirs(cur_save_dir, exist_ok=True)
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    # 以时间格式命名
    video_path = osp.join(cur_save_dir, "{}.mp4".format(ti))

    # 视频流的宽高
    width = cap.get(3)
    height = cap.get(4)
    fps = 20
    writer = cv2.VideoWriter(video_path, fourcc, fps, (int(width), int(height)))
    frame_cnt = 0
    # 需要的存储的帧数
    while (frame_cnt < save_time * fps):
        frame_cnt += 1
        ret,frame = cap.read()
        if not ret:
            print("frame is empty")
            break
        writer.write(frame)
    cap.release()
    writer.release()

    return ret

def video_capture():
    # cap = cv2.VideoCapture(0)
    cap = cv2.VideoCapture("rtmp://140.246.93.130:1935/vod/openUrl/mBLVC3C?beginTime=20211205T174929&endTime=20211205T175029&playBackMode=1")
    # cap = cv2.VideoCapture("rtsp://admin:abcd1234@192.168.104.101:554/ch01/sub/av_stream")

    #    const std::string videoStreamAddress = "rtsp://admin:swlskjyxgs0617@192.168.31.111:554/h264/ch34/main/av_stream";
    print("s")
    ret,frame = cap.read()
    while ret:
        ret,frame = cap.read()
        print(frame.shape)
        cv2.imshow("frame",frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cv2.destroyAllWindows()
    cap.release()

def grab_video_by_time(cap_ips, save_dir, per_minute=10):
    his_data = {}

    start = time.time()
    """确定起始时间，不断叠加时间"""
    while True:
        cur_time = time.time() 
        time_diff = cur_time - start
        if time_diff < per_minute * 60:
            print(time_diff)
            time.sleep(1)
            continue

        for ci in cap_ips:
            # his_data中有没有取到这些数据
            if his_data.get(ci[0], None) and \
                his_data[ci[0]][0] == False and \
                cur_time - his_data[ci[0]][1] < 30 * 60:
                continue
            # ret = save_img_by_rtsp(save_dir, ci)
            ret = save_video_by_rtsp(save_dir, ci, 10)
            his_data[ci[0]] = [ret, cur_time]

        start = time.time()

if __name__ == "__main__":

    cap_ips = [
        #       ip            usr       pwd
        ["192.168.104.101", "admin", "abcd1234"],
        ["192.168.104.102", "admin", "abcd1234"],
        ["192.168.104.104", "admin", "abcd1234"],
        ["192.168.104.105", "admin", "abcd1234"],
        ["192.168.104.107", "admin", "abcd1234"],
        ["192.168.104.108", "admin", "abcd1234"],
        ["192.168.104.110", "admin", "abcd1234"],
        ["192.168.104.111", "admin", "abcd1234"],
        ["192.168.104.113", "admin", "abcd1234"],
        ["192.168.104.114", "admin", "abcd1234"],
        ["192.168.104.116", "admin", "abcd1234"],
        ["192.168.104.117", "admin", "abcd1234"],
        ["192.168.104.118", "admin", "abcd1234"],
        ["192.168.104.119", "admin", "abcd1234"],

        ["192.168.107.121", "admin", "abcd1234"],
        ["192.168.107.122", "admin", "abcd1234"],
        ["192.168.107.124", "admin", "abcd1234"],
        ["192.168.107.125", "admin", "abcd1234"],
        ["192.168.107.127", "admin", "abcd1234"],
        ["192.168.107.128", "admin", "abcd1234"],
        ["192.168.107.130", "admin", "abcd1234"],
        ["192.168.107.131", "admin", "abcd1234"],
    ]

    # orginal
    save_dir = r"F:\work\dataset\video\leishan"
    # save_dir = r"E:\gec\sample\sampledata\video"
    per_minute = 20
    grab_video_by_time(cap_ips, save_dir, per_minute)

    # # 雷山
    # cap_ips = [
    #     #       ip            usr       psw   channel
    #     ["192.168.104.103", "admin", "abcd1234", 2] ,
    #     ["192.168.104.106", "admin", "abcd1234", 2] ,
    #     ["192.168.104.109", "admin", "abcd1234", 2] ,
    #     ["192.168.104.112", "admin", "abcd1234", 2] ,
    #     ["192.168.104.115", "admin", "abcd1234", 2] ,
    #     ["192.168.104.120", "admin", "abcd1234", 2] ,
    #     ["192.168.107.123", "admin", "abcd1234", 2] ,
    #     ["192.168.107.126", "admin", "abcd1234", 2] ,
    #     ["192.168.107.129", "admin", "abcd1234", 2] ,
    #     ["192.168.107.132", "admin", "abcd1234", 2] ,
    # ]
    # grab_list = ["2021-12-22 00:30:00", "2021-12-29 23:30:00", 1 * 60 * 60]
    # save_dir = r"F:\work\code\data\video\leishan"
    # for cap_ip in cap_ips:
    #     save_video_by_rtsp_history(save_dir, cap_ip, grab_list[0], grab_list[1], grab_list[2])

    # # 长顺
    # cap_ips = [
    #     #       ip            usr       psw   channel
    #     # ["192.168.100.168", "admin", "abcd1234", 1] ,
    #     # ["192.168.101.168", "admin", "abcd1234", 1] ,
    #     # ["192.168.102.168", "admin", "abcd1234", 1] ,
    #     ["192.168.103.168", "admin", "abcd1234", 1] ,
    #     ["192.168.104.168", "admin", "abcd1234", 1] ,
    #     ["192.168.105.168", "admin", "abcd1234", 1] ,
    #     ["192.168.106.168", "admin", "abcd1234", 1] ,
    #     ["192.168.107.168", "admin", "abcd1234", 1] ,
    #     # ["192.168.108.168", "admin", "abcd1234", 1] ,
    #     # ["192.168.109.168", "admin", "abcd1234", 1] ,
    #     # ["192.168.110.168", "admin", "abcd1234", 1] ,
    #     # ["192.168.111.168", "admin", "abcd1234", 1] ,
    #     # ["192.168.112.168", "admin", "abcd1234", 1] ,
    # ]
    # # grab_list = ["2021-11-01 00:30:00", "2021-12-30 23:30:00", 15 * 60]
    # # # save_dir = r"F:\work\dataset\pic\changshun\car"
    # # save_dir = r"E:\gec\sample\sampledata\video"
    # # for cap_ip in cap_ips:
    # #     save_img_by_rtsp_history(save_dir, cap_ip, grab_list[0], grab_list[1], grab_list[2])

  